#!/usr/bin/python

import xlh
import sys, getopt, os
from random import Random

# Define some general functions


def exitOnError(errmsg, exit):
	print 'The following error(s) occurred:\n'
	print errmsg
	sys.exit(exit)
def usage():
	print """Usage:

xdbcreate.py [-h] -d DATABASE [-f FILE | -a ACCESSIONS] -m MISSED [-e ENZYME] [-o OUTPUT] [-l LENGTH]

 -h:	Print this help and exit
 -d:	Database used for previous MS search
 -f:	File containing accession numbers
 -a:	Comma separated accession numbers
 -m:	Number of missed cleavages used in previous MS search [Default 2]
 -e:	Enzyme [Default: Trypsin]
 -o:	Name of output database [Optional]
 -l:	Minumum peptide length [Default 3]
 -L:	Maximum peptide length [Default 15]
 -A:	Use all peptides (not only xlinking ones) [Default False]
	"""
	sys.exit(1)

def readAccessionFile(fileName):
	accs = []
	fh = open(fileName, 'r')
	while 1:
		l = fh.readline()
		if not l:
			break
		accs.append(l.strip())
	fh.close()
	return accs
	

MinLen = 3
MaxLen = 15
DbIn = ''
Accs = []
Missed = 2
DbOut = ''
Function = 0
Enzyme = 'Trypsin'
UseAll = False

try:
	optlist, args = getopt.getopt(sys.argv[1:], 'd:f:a:m:o:e:l:L:A')
except getopt.GetoptError:
	usage()

if not optlist:
	usage()

for o, a in optlist:
	if o == '-h':
		usage()
	if o == '-d':
		DbIn = a
	if o == '-o':
		DbOut = a
	if o == '-m':
		Missed = int(a)
	if o == '-l':
		MinLen = int(a)
	if o == '-L':
		MaxLen = int(a)
	if o == '-f':
		Accs = readAccessionFile(a)
	if o == '-a':
		Accs = a.split(',')
	if o == '-e':
		Enzyme = a.strip()
	if o == '-A':
		UseAll = True
		
# do some checks
if not DbIn:
	exitOnError('Missing database name', 2)
if not Accs:
	exitOnError('No accession numbers defined', 2)
if not DbOut:
	DbOut = 'db' + str(Random().getrandbits(64)) 

FastaOut = DbOut + '.fasta'
DBtitle = 'Used ' + Enzyme + ' with ' + str(Missed) + ' miscleavages'

# read proteins from database

SingleProteins = []
for acc in Accs:
	SingleProteins.append(xlh.Protein(xlh.readEntryFromNCBI(DbIn, acc)))
	SingleProteins[-1].cleave(Enzyme, Missed)
	SingleProteins[-1]._peptides = filter(lambda x: len(x.sequence()) >= MinLen and len(x.sequence()) <= MaxLen, SingleProteins[-1].peptides())
	SingleProteins[-1]._peptides = filter(lambda x:  'K' in x.sequence()[:-1] or x.isNTerminal(SingleProteins[-1]), SingleProteins[-1].peptides())
	
# in this case we should create a pool of peptides
# and concatenate them ABA, BCB, CDC, DAD etc.
# each concatenation is a single fasta entry
t = {}
entries = []
for pr in SingleProteins:
	peptides = pr.peptideSequences()
	for pt in peptides:
		try:
			t[pt].append(pr.accession())
		except KeyError:
			t[pt] = [pr.accession()]

#sys.stdout.write("Got " + str(len(peps)) + " peptides\n")


linewrap = 80
thisDir = os.getcwd()
if '/' not in FastaOut:
	# if no directory is given, assume that is the
	# one specified in external.ini
	os.chdir(xlh.External.get("blastdb", "path"))

n = 0
for px in range(len(SingleProteins) -1 ):
	for py in range(px, len(SingleProteins)):
		Name = FastaOut + '.part' + str(n)
		Db = DbOut + '.part' + str(n)
		if len(SingleProteins[px].peptideSequences()) == 0 or len(SingleProteins[py].peptideSequences()) == 0:
			continue
		sys.stdout.write("Writing " + Name)
		sys.stdout.flush()
		fh = open(Name, 'w')
		for x in range(len(SingleProteins[px].peptideSequences())):
			for y in range(len(SingleProteins[py].peptideSequences())):
				acc = 'xpep' + str(x) + 'x' + str(y)
				desc = 'xlh|' + acc + '| ' + SingleProteins[px].peptideSequences()[x] + ':' + ','.join(t[SingleProteins[px].peptideSequences()[x]]) + ' ' + SingleProteins[py].peptideSequences()[y] + ':'  + ','.join(t[SingleProteins[py].peptideSequences()[y]])
				seq = SingleProteins[px].peptideSequences()[x] + SingleProteins[py].peptideSequences()[y] + SingleProteins[px].peptideSequences()[x]
				fh.write('>' + desc + '\n')
				for i in range(len(seq) / linewrap):
					fh.write(seq[i * linewrap:i * linewrap + linewrap] + '\n')
				r = len(seq) % linewrap
				fh.write(seq[-r:] + '\n')
		n += 1
		fh.close()


		sys.stdout.flush()
		sys.stdout.write(" OK\n")

		sys.stdout.write("Formatting for omssa search...")
		sys.stdout.flush()
		xlh.writeNCBI(Name, Db, DBtitle)
		sys.stdout.write(" OK\n")
os.chdir(thisDir)

